#!/usr/bin/python3
# -*- coding: utf-8 -*-

"""
SQLAlchemy工具模块

SQLAlchemy ORM集成，提供数据库模型基类和管理功能
"""

import os
from contextlib import contextmanager
from datetime import datetime
from typing import Any, Dict, Optional, Type, TypeVar
from sqlalchemy import create_engine, Column, String, DateTime, inspect
from sqlalchemy.orm import declarative_base, sessionmaker, scoped_session
from sqlalchemy.pool import QueuePool
from flask import Flask, current_app, g

# 创建模型基类
Base = declarative_base()

# 定义模型类型变量
T = TypeVar('T', bound=Base)


class BaseModel(Base):
    """ORM模型基类，包含共同的字段和方法"""
    
    __abstract__ = True  # 标记为抽象类，不创建表
    
    id = Column(String(50), primary_key=True, doc="主键ID")
    created_time = Column(DateTime, default=datetime.now, doc="创建时间")
    updated_time = Column(DateTime, default=datetime.now, onupdate=datetime.now, doc="更新时间")
    
    def to_dict(self, exclude: Optional[list] = None) -> Dict[str, Any]:
        """将模型转换为字典
        
        Args:
            exclude: 要排除的字段列表
            
        Returns:
            Dict[str, Any]: 模型数据字典
        """
        exclude = exclude or []
        return {
            c.key: getattr(self, c.key)
            for c in inspect(self).mapper.column_attrs
            if c.key not in exclude
        }
    
    @classmethod
    def from_dict(cls: Type[T], data: Dict[str, Any]) -> T:
        """从字典创建模型实例
        
        Args:
            data: 模型数据字典
            
        Returns:
            T: 模型实例
        """
        return cls(**{
            k: v for k, v in data.items()
            if k in inspect(cls).columns.keys()
        })


def init_sqlalchemy(app: Flask):
    """初始化SQLAlchemy
    
    Args:
        app: Flask应用实例
    """
    # 获取数据库路径
    db_path = app.config['DB_PATH']
    db_uri = f"sqlite:///{db_path}"
    
    # 创建引擎
    engine = create_engine(
        db_uri,
        # SQLite连接池配置
        poolclass=QueuePool,
        pool_size=5,
        pool_timeout=30,
        pool_recycle=3600,
        # 启用外键支持
        connect_args={"check_same_thread": False},
    )
    
    # 创建会话工厂
    session_factory = sessionmaker(bind=engine)
    
    # 创建线程安全的会话
    session = scoped_session(session_factory)
    
    # 将会话绑定到模型元数据
    Base.query = session.query_property()
    
    # 将SQLAlchemy对象保存到应用上下文
    app.extensions = getattr(app, 'extensions', {})
    app.extensions['sqlalchemy'] = {
        'engine': engine,
        'session': session,
        'Base': Base
    }
    
    @app.teardown_appcontext
    def cleanup_session(exception=None):
        """关闭会话"""
        if 'sqlalchemy_session' in g:
            g.sqlalchemy_session.close()


@contextmanager
def get_session():
    """SQLAlchemy会话上下文管理器
    
    使用示例:
    ```
    with get_session() as session:
        users = session.query(User).all()
    ```
    
    Yields:
        sqlalchemy.orm.Session: 数据库会话
    """
    # 获取SQLAlchemy会话
    if 'sqlalchemy_session' not in g:
        g.sqlalchemy_session = current_app.extensions['sqlalchemy']['session']()
    
    session = g.sqlalchemy_session
    
    try:
        # 提供会话
        yield session
        # 没有异常则提交
        session.commit()
    except Exception as e:
        # 出现异常回滚
        session.rollback()
        raise e


class SQLAlchemyRepository:
    """SQLAlchemy通用仓储类"""
    
    def __init__(self, model_class: Type[T]):
        """初始化仓储
        
        Args:
            model_class: 模型类
        """
        self.model_class = model_class
    
    def get_by_id(self, id: str) -> Optional[T]:
        """根据ID获取实体
        
        Args:
            id: 实体ID
            
        Returns:
            Optional[T]: 实体对象或None
        """
        with get_session() as session:
            return session.query(self.model_class).filter_by(id=id).first()
    
    def list(self, filters: Dict[str, Any] = None, page: int = 1, size: int = 10) -> tuple:
        """获取实体列表
        
        Args:
            filters: 过滤条件
            page: 页码
            size: 每页大小
            
        Returns:
            tuple: (实体列表, 总数)
        """
        filters = filters or {}
        
        with get_session() as session:
            query = session.query(self.model_class)
            
            # 应用过滤条件
            for key, value in filters.items():
                if hasattr(self.model_class, key):
                    query = query.filter(getattr(self.model_class, key) == value)
            
            # 获取总数
            total = query.count()
            
            # 分页
            offset = (page - 1) * size
            entities = query.limit(size).offset(offset).all()
            
            return entities, total
    
    def create(self, entity: T) -> T:
        """创建实体
        
        Args:
            entity: 实体对象
            
        Returns:
            T: 创建后的实体
        """
        with get_session() as session:
            session.add(entity)
            session.flush()  # 刷新以获取ID
            session.refresh(entity)  # 刷新实体
            return entity
    
    def update(self, entity: T) -> T:
        """更新实体
        
        Args:
            entity: 实体对象
            
        Returns:
            T: 更新后的实体
        """
        with get_session() as session:
            session.merge(entity)
            session.flush()
            session.refresh(entity)
            return entity
    
    def delete(self, id: str) -> bool:
        """删除实体
        
        Args:
            id: 实体ID
            
        Returns:
            bool: 是否成功删除
        """
        with get_session() as session:
            entity = session.query(self.model_class).filter_by(id=id).first()
            if entity:
                session.delete(entity)
                return True
            return False 